import onnx
import onnx_graphsurgeon as gs

# 加载模型
graph = gs.import_onnx(onnx.load("original_model.onnx"))


#删除 conv2
def remove_add_node():
    # 找到 Add 节点
    conv1_node = next(node for node in graph.nodes if node.name == "/conv1/Conv")
    conv2_node = next(node for node in graph.nodes if node.name == "/conv2/Conv")
    conv3_node = next(node for node in graph.nodes if node.name == "/conv3/Conv")

    print("\n=== 删除前 ===")
    print(f"找到待删除节点: {conv2_node.name}")
    print("当前节点列表：", [node.name for node in graph.nodes])

    # 执行删除
    graph.nodes.remove(conv2_node)
    conv3_node.inputs[0] = conv1_node.outputs[0]
    
    graph.cleanup()  # 必须清理孤立张量

    print("\n=== 删除后 ===")
    print("更新后节点列表：", [node.name for node in graph.nodes])

    print("\n=== 保存中 ===")
    # 保存修改后的模型
    onnx.save(gs.export_onnx(graph), "removed_add.onnx")

#删除 Add
def remove_add_node1():
    # 找到 Add 节点
    relu_node = next(node for node in graph.nodes if node.name == "/relu/Relu")
    add_node = next(node for node in graph.nodes if node.name == "/Add")

    print("\n=== 删除前 ===")
    print(f"找到待删除节点: {add_node.name}")
    print("当前节点列表：", [node.name for node in graph.nodes])

    # 执行删除
    graph.nodes.remove(add_node)
    graph.outputs[0] = add_node.inputs[0] # graph.outputs[0] = relu_node.outputs[0] # 这样也可
    # # graph.outputs[0].dtype = onnx.TensorProto.FLOAT  # 设置为float32类型
    graph.outputs[0].dtype = add_node.outputs[0].dtype
    graph.outputs[0].shape = add_node.outputs[0].shape
    graph.outputs[0].name = "output"

    # 还可以如下
    # graph.outputs = [add_node.inputs[0]]
    # graph.outputs[0].dtype = add_node.outputs[0].dtype
    # # graph.outputs[0].shape = add_node.outputs[0].shape
    # graph.outputs[0].name = "output"



    graph.cleanup()  # 必须清理孤立张量

    print("\n=== 删除后 ===")
    print("更新后节点列表：", [node.name for node in graph.nodes])

    print("\n=== 保存中 ===")
    # 保存修改后的模型
    onnx.save(gs.export_onnx(graph), "removed_add.onnx")

if __name__ == "__main__":
    remove_add_node1()